{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Recommender Systems with DGL\n",
    "\n",
    "## Introduction\n",
    "\n",
    "Graph Neural Networks (GNN), as a methodology of learning representations on graphs, has gained much attention recently.  Various models such as Graph Convolutional Networks, GraphSAGE, etc. are proposed to obtain representations of whole graphs, or nodes on a single graph.\n",
    "\n",
    "A primary goal of recommendation is to automatically make predictions about a user's interest, e.g. whether/how a user would interact with a set of items, given the interaction history of the user herself, as well as the histories of other users.  The user-item interaction can also be viewed as a bipartite graph, where users and items form two sets of nodes, and edges connecting them stands for interactions.  The problem can then be formulated as a *link-prediction* problem, where we try to predict whether an edge (of a given type) exists between two nodes.\n",
    "\n",
    "Based on this intuition, the academia developed multiple new models for recommendation, including but not limited to:\n",
    "\n",
    "* Geometric Learning Approaches\n",
    "  * [Geometric Matrix Completion](https://papers.nips.cc/paper/5938-collaborative-filtering-with-graph-information-consistency-and-scalable-methods.pdf)\n",
    "  * [Recurrent Multi-graph CNN](https://arxiv.org/pdf/1704.06803.pdf)\n",
    "* Graph-convolutional Approaches\n",
    "  * Models such as [R-GCN](https://arxiv.org/pdf/1703.06103.pdf) or [GraphSAGE](https://github.com/stellargraph/stellargraph/tree/develop/demos/link-prediction/hinsage) also apply.\n",
    "  * [Graph Convolutional Matrix Completion](https://arxiv.org/abs/1706.02263)\n",
    "  * [PinSage](https://arxiv.org/pdf/1806.01973.pdf)\n",
    "  \n",
    "In this hands-on tutorial, we will demonstrate how to write a recommendation model with GraphSAGE in DGL + MXNet. We just demonstrate rating prediction in this tutorial."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import dgl\n",
    "import dgl.function as FN\n",
    "\n",
    "# Load MXNet as backend\n",
    "dgl.load_backend('mxnet')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import mxnet as mx\n",
    "from mxnet import ndarray as nd, autograd, gluon\n",
    "from mxnet.gluon import nn\n",
    "import numpy as np\n",
    "\n",
    "np.random.seed(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading data\n",
    "\n",
    "In this tutorial, we focus on rating prediction on MovieLens-100K dataset.  The data comes from [MovieLens](http://files.grouplens.org/datasets/movielens/ml-1m.zip) and is shipped with the notebook already.\n",
    "\n",
    "The dataset is encapsulated in `movielens.MovieLens` class for clarity of this notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import movielens\n",
    "\n",
    "ml = movielens.MovieLens('ml-100k')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The MovieLens graph has 2,625 nodes and 200,000 edges. The graph contains both user nodes and movie nodes. The MovieLens dataset contains 100,000 ratings and the ratings are stored as edges. DGL only supports directed edges. Thus, each rating is stored twice: one connects from a user node to a movie node, and the other connects from a movie node to a user node."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This is a DGLGraph object.\n",
    "g = ml.g\n",
    "print('#vertices:', g.number_of_nodes())\n",
    "print('#edges:', g.number_of_edges())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## See the features in the MovieLens dataset\n",
    "\n",
    "The MovieLens dataset has some user features and movie features.\n",
    "\n",
    "User features:\n",
    "* age,\n",
    "* gender,\n",
    "* occupation,\n",
    "* zip code,\n",
    "\n",
    "Movie features:\n",
    "* genre,\n",
    "* year,\n",
    "\n",
    "We use one-hot encoding for \"age\", \"gender\", \"occupation\", \"zip code\" and \"year\". \"genre\" uses multi-hot encoding. We then store them as node features on the graph.\n",
    "\n",
    "Since user features and item features are different, we pad both types of features with zeros."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(g.ndata)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There is a node data \"type\" that indicates the node type in the bipartite graph. Nodes with \"type=1\" are user nodes and \"type=0\" are movie nodes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('#user nodes:', mx.nd.sum(g.ndata['type'] == 1).asnumpy())\n",
    "print('#movie nodes:', mx.nd.sum(g.ndata['type'] == 0).asnumpy())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Recommendation model with GNN\n",
    "\n",
    "<img src=\"https://s3.us-east-2.amazonaws.com/dgl.ai/amlc_tutorial/rec_process.png\" width=\"800\">\n",
    "\n",
    "Recommendation with graph neural networks has two steps:\n",
    "* graph encoder: use graph neural networks to compute node embeddings.\n",
    "* edge decoder: compute scores on edges with user embeddings and movie embeddings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GNNRecommender(nn.Block):\n",
    "    def __init__(self, encoder, decoder):\n",
    "        super(GNNRecommender, self).__init__()\n",
    "        \n",
    "        self.encoder = encoder\n",
    "        self.decoder = decoder\n",
    "\n",
    "    def forward(self, G, users, items):\n",
    "        h = self.encoder(G)\n",
    "        h_users = h[users]\n",
    "        h_items = h[items]\n",
    "        return self.decoder(h_users, h_items, users, items)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## GraphSage encoder"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The encoder does two things:\n",
    "* generate the initial user and movie embeddings from categorial features.\n",
    "* run GraphSAGE layers on the MovieLens graph multiple times to compute the final embeddings for rating prediction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GraphSageEncoder(nn.Block):\n",
    "    def __init__(self, embedding_size, n_layers, G):\n",
    "        super(GraphSageEncoder, self).__init__()\n",
    "\n",
    "        self.G = G\n",
    "        self.user_nodes = G.filter_nodes(lambda nodes: nodes.data['type'] == 1).astype(np.int64)\n",
    "        self.movie_nodes = G.filter_nodes(lambda nodes: nodes.data['type'] == 0).astype(np.int64)\n",
    "\n",
    "        self.layers = nn.Sequential()\n",
    "        for i in range(n_layers):\n",
    "            self.layers.add(GraphSageLayer(embedding_size, G, self.user_nodes, self.movie_nodes))\n",
    "\n",
    "        # One-hot encoding for each node.\n",
    "        node_emb = nn.Embedding(G.number_of_nodes() + 1, embedding_size)\n",
    "        self.user_emb = UserEmbedding(G, embedding_size, node_emb)\n",
    "        self.movie_emb = MovieEmbedding(G, embedding_size, node_emb)\n",
    "\n",
    "    def forward(self, G):\n",
    "        # Generate embeddings on user nodes and movie nodes.\n",
    "        assert G.number_of_nodes() == self.G.number_of_nodes()\n",
    "        G.apply_nodes(lambda nodes: {'h': self.user_emb(nodes.data, self.user_nodes)}, self.user_nodes)\n",
    "        G.apply_nodes(lambda nodes: {'h': self.movie_emb(nodes.data, self.movie_nodes)}, self.movie_nodes)\n",
    "\n",
    "        for layer in self.layers:\n",
    "            layer(G)\n",
    "\n",
    "        # The node embeddings computed by GraphSage layers are stored in node data of 'h'.\n",
    "        return G.ndata['h']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## GraphSage model\n",
    "\n",
    "We can now write a GraphSAGE layer.  In GraphSAGE, the node representation is updated with the representation in the previous layer as well as an aggregation (often mean) of \"messages\" sent from all neighboring nodes.\n",
    "\n",
    "### Algorithm\n",
    "\n",
    "The algorithm of a single GraphSAGE layer has three steps for each node $v$:\n",
    "\n",
    "1. $h_{\\mathcal{N}(v)} \\gets \\mathtt{Sum}_{u \\in \\mathcal{N}(v)} h_{u}$\n",
    "2. $h_{v} \\gets \\sigma\\left(W \\cdot \\mathtt{CONCAT}(h_v, h_{\\mathcal{N}(v)}/d_{\\mathcal{N}(v)})\\right)$\n",
    "3. $h_{v} \\gets h_{v} / \\lVert h_{v} \\rVert_2$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Slight modification on the original GraphSage model\n",
    "\n",
    "In practice, the MovieLens dataset has two types of nodes: users and movies. We need to perform separate node update functions on the two types of nodes on step 2.\n",
    "\n",
    "For the movie nodes,\n",
    "\n",
    "2. $h_{m} \\gets \\sigma\\left(W0 \\cdot \\mathtt{CONCAT}(h_m, h_{\\mathcal{N}(m)} / d_{\\mathcal{N}(m)})\\right)$\n",
    "\n",
    "For the user nodes,\n",
    "\n",
    "2. $h_{u} \\gets \\sigma\\left(W1 \\cdot \\mathtt{CONCAT}(h_u, h_{\\mathcal{N}(u)} / d_{\\mathcal{N}(u)})\\right)$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GraphSageLayer(nn.Block):\n",
    "    def __init__(self, feature_size, G, user_nodes, movie_nodes):\n",
    "        super(GraphSageLayer, self).__init__()\n",
    "\n",
    "        self.feature_size = feature_size\n",
    "        self.G = G\n",
    "        self.user_nodes = user_nodes\n",
    "        self.movie_nodes = movie_nodes\n",
    "\n",
    "        self.user_update = GraphSageNodeUpdate(feature_size)\n",
    "        self.movie_update = GraphSageNodeUpdate(feature_size)\n",
    "\n",
    "        all_nodes = mx.nd.arange(G.number_of_nodes(), dtype=np.int64)\n",
    "        self.deg = G.in_degrees(all_nodes).astype(np.float32)\n",
    "\n",
    "    def forward(self, G):\n",
    "        assert G.number_of_nodes() == self.G.number_of_nodes()\n",
    "        G.ndata['deg'] = self.deg\n",
    "        # Step 1\n",
    "        # The message function and message reduce function used by GraphSage.\n",
    "        # EXERCISE: Interested users can try different message functions and reduce functions.\n",
    "        G.update_all(FN.copy_src('h', 'h'), FN.sum('h', 'h_agg'))\n",
    "        # Step 2 and 3\n",
    "        G.apply_nodes(self.user_update, self.user_nodes)\n",
    "        G.apply_nodes(self.movie_update, self.movie_nodes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GraphSageNodeUpdate(nn.Block):\n",
    "    def __init__(self, feature_size):\n",
    "        super(GraphSageNodeUpdate, self).__init__()\n",
    "\n",
    "        self.feature_size = feature_size\n",
    "        self.W = nn.Dense(feature_size)\n",
    "        self.leaky_relu = nn.LeakyReLU(0.1)\n",
    "\n",
    "    def forward(self, nodes):\n",
    "        # Node embedding from the previous layer.\n",
    "        h = nodes.data['h']\n",
    "        # Aggregation of the node embeddings in the neighborhood\n",
    "        h_agg = nodes.data['h_agg']\n",
    "        # Degree of the vertex.\n",
    "        deg = nodes.data['deg'].expand_dims(1)\n",
    "        h_concat = nd.concat(h, h_agg / nd.maximum(deg, 1e-6), dim=1)\n",
    "        h_new = self.leaky_relu(self.W(h_concat))\n",
    "        # Layer norm\n",
    "        return {'h': h_new / nd.maximum(h_new.norm(axis=1, keepdims=True), 1e-6)}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compute node embeddings on the MovieLens dataset\n",
    "\n",
    "User nodes and movie nodes have different sets of features. Thus, we need to generate embeddings differently.\n",
    "\n",
    "User nodes have categorial features of \"age\", \"gender\", \"occupation\" and \"zip code\". These features are all one-hot encodings. In addition, we add one-hot encoding for every user node. To generate user embedding, we add all of these embeddings together."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class UserEmbedding(nn.Block):\n",
    "    def __init__(self, G, feature_size, node_emb):\n",
    "        super(UserEmbedding, self).__init__()\n",
    "\n",
    "        # Embedding matrices for one-hot encoding.\n",
    "        self.emb_age = nn.Embedding(G.ndata['age'].max().asscalar() + 1,\n",
    "                                    feature_size)\n",
    "        self.emb_gender = nn.Embedding(G.ndata['gender'].max().asscalar() + 1,\n",
    "                                       feature_size)\n",
    "        self.emb_occupation = nn.Embedding(G.ndata['occupation'].max().asscalar() + 1,\n",
    "                                           feature_size)\n",
    "        self.emb_zip = nn.Embedding(G.ndata['zip'].max().asscalar() + 1,\n",
    "                                    feature_size)\n",
    "        \n",
    "        # One-hot encoding for each node.\n",
    "        self.node_emb = node_emb\n",
    "\n",
    "    def forward(self, ndata, nid):\n",
    "        h = self.node_emb(nid + 1)\n",
    "        extra_repr = []\n",
    "        extra_repr.append(self.emb_age(ndata['age']))\n",
    "        extra_repr.append(self.emb_gender(ndata['gender']))\n",
    "        extra_repr.append(self.emb_occupation(ndata['occupation']))\n",
    "        extra_repr.append(self.emb_zip(ndata['zip']))\n",
    "        return h + nd.stack(*extra_repr, axis=0).sum(axis=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Movie nodes \"year\", \"genre\". \"year\" is one-hot encoding, \"genre\" is stored in a float32 dense matrix. Like user nodes, we add one-hot encoding to every movie node. To generate movie embedding, we add all of these embeddings together."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MovieEmbedding(nn.Block):\n",
    "    def __init__(self, G, feature_size, node_emb):\n",
    "        super(MovieEmbedding, self).__init__()\n",
    "        self.emb_year = nn.Embedding(G.ndata['year'].max().asscalar() + 1,\n",
    "                                     feature_size)\n",
    "\n",
    "        # Linear projection for float32 features.\n",
    "        seq = nn.Sequential()\n",
    "        with seq.name_scope():\n",
    "            seq.add(nn.Dense(feature_size))\n",
    "            seq.add(nn.LeakyReLU(0.1))\n",
    "        self.proj_genre = seq\n",
    "\n",
    "        # One-hot encoding for each node.\n",
    "        self.node_emb = node_emb\n",
    "\n",
    "    def forward(self, ndata, nid):\n",
    "        h = self.node_emb(nid + 1)\n",
    "        extra_repr = []\n",
    "        extra_repr.append(self.emb_year(ndata['year']))\n",
    "        extra_repr.append(self.proj_genre(ndata['genre']))\n",
    "        return h + nd.stack(*extra_repr, axis=0).sum(axis=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Rating prediction\n",
    "\n",
    "For recommendation, the rating on item $j$ by user $i$ is defined by $u_i^T v_j$.\n",
    "\n",
    "In practice, recommendation models have user bias term and movie bias term: $u_i^T v_j + b_{u_i} + b_{v_j}$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DotDecoder(nn.Block):\n",
    "    def __init__(self, num_nodes):\n",
    "        super(DotDecoder, self).__init__()\n",
    "        \n",
    "        with self.name_scope():\n",
    "            self.biases = self.params.get(\n",
    "                'node_biases',\n",
    "                init=mx.init.Zero(),\n",
    "                shape=(num_nodes+1,))\n",
    "            \n",
    "    def forward(self, h_users, h_items, users, items):\n",
    "        user_biases = self.biases.data()[users+1]\n",
    "        item_biases = self.biases.data()[items+1]\n",
    "        return (h_users * h_items).sum(1) + user_biases + item_biases"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Get the training set\n",
    "\n",
    "We use 80% of edges for training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "g_train = ml.g_train\n",
    "rating_train = g_train.edata['rating']\n",
    "src_train, dst_train = g_train.all_edges()\n",
    "print('#vertices:', g_train.number_of_nodes())\n",
    "print('#edges:', g_train.number_of_edges())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Get the testing edge set\n",
    "\n",
    "We use 20% of edges for evaluation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "g = ml.g\n",
    "eid_test = g.filter_edges(lambda edges: edges.data['test']).astype('int64')\n",
    "src_test, dst_test = g.find_edges(eid_test)\n",
    "rating_test = g.edges[eid_test].data['rating']\n",
    "print('#testing edges:', len(eid_test))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# We use 1-layer GraphSage as the graph encoder.\n",
    "# EXERCISE: Interested users can try different numbers of layers (e.g., 0, 1, 2)\n",
    "model = GNNRecommender(GraphSageEncoder(embedding_size=100, n_layers=1, G=g_train),\n",
    "                       DotDecoder(g_train.number_of_nodes()))\n",
    "model.collect_params().initialize(ctx=mx.cpu())\n",
    "trainer = gluon.Trainer(model.collect_params(), 'adam', {'learning_rate': 0.003, 'wd': 1e-5})\n",
    "\n",
    "for epoch in range(200):\n",
    "    # Training\n",
    "    for _ in range(10):\n",
    "        with mx.autograd.record():\n",
    "            score = model(g_train, src_train, dst_train)\n",
    "            loss = ((score - rating_train) ** 2).mean()\n",
    "            loss.backward()\n",
    "        trainer.step(1)\n",
    "\n",
    "    # Testing\n",
    "    h = model.encoder(g)\n",
    "    # Compute test RMSE\n",
    "    score = model.decoder(h[src_test], h[dst_test], src_test, dst_test)\n",
    "    test_rmse = nd.sqrt(((score - rating_test) ** 2).mean())\n",
    "\n",
    "    print('Training loss:', loss.asscalar(), 'Test RMSE:', test_rmse.asscalar())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if 'h' in g_train.ndata:\n",
    "    del g_train.ndata['h']\n",
    "if 'h_agg' in g_train.ndata:\n",
    "    del g_train.ndata['h_agg']\n",
    "if 'h' in g.ndata:\n",
    "    del g.ndata['h']\n",
    "if 'h_agg' in g.ndata:\n",
    "    del g.ndata['h_agg']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Take-home Exercise\n",
    "\n",
    "Interested users can try different things given this simple recommendation model with GraphSage. Interesting experiments include:\n",
    "* try different numbers of GraphSage layers.\n",
    "* try different graph encoders. For example, replace the message passing functions with the one in [GAT](https://github.com/dmlc/dgl/blob/master/examples/mxnet/gat/gat.py) or GCMC.\n",
    "* try mini-batch training. A tutorial of training GraphSage with mini-batch has been included in the same folder of this tutorial.\n",
    "* try different recommendation datasets. To explore the benefit of graph neural networks, we suggest users to try feature-rich datasets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_python3",
   "language": "python",
   "name": "conda_python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
